from copy import deepcopy
from dataclasses import dataclass, field
import sys
from typing import Literal

from PyBirdViewCode.uast.uast_commands import extract_cfg_from_method
from PyBirdViewCode.uast.universal_code_property_graphs import (
    compose_pdg_topology,
    get_cdg_topology,
    add_label,
    CodePropertyGraphs,
    get_ddg_topology,
)

sys.path.append("../../")
import networkx as nx

# from PyBirdViewCode
from PyBirdViewCode import (
    FileManager,
    get_file_uast,
    get_method_cpg,
    get_pdg_from_cfg,
    build_control_dependence_graph,
    # add_label,
    UASTQuery,
    uast,
    BaseUASTUnparser,
)

# 在当前路径下创建FileManager对象
fm = FileManager("./output")


def set_edge_label(g: nx.DiGraph, **attrs: dict):
    """
    将cfg_nx的边标签添加到g上
    """
    for u, v in g.edges:
        # if u != "ENTRY":
        for attr, value in attrs.items():
            g.edges[u, v][attr] = value


class Analyzer:
    def __init__(self) -> None:
        # 获取文件的UAST
        self.file_uast = get_file_uast("demo.c")
        self.unparser = BaseUASTUnparser()

    def create_input_nodes(self, funcdef_uast: uast.MethodDecl):

        init_vertices = []
        final_vertices = []

        # 解析函数定义
        for formal_param in funcdef_uast.type.pos_args:
            match formal_param:
                case uast.ParamDecl(type=uast.AddrReferenceType(), name=arg_name):
                    init_vertices.append(
                        uast.Assignment(
                            "=",
                            lhs=[uast.DereferenceExpr(arg_name)],
                            rhs=[
                                uast.DereferenceExpr(
                                    uast.Name(
                                        f"$delta<{self.unparser.unparse(funcdef_uast.name)}>({self.unparser.unparse(arg_name)})"
                                    )
                                )
                            ],
                        )
                    )
                    final_vertices.append(
                        uast.Assignment(
                            "=",
                            lhs=[
                                uast.DereferenceExpr(
                                    uast.Name(
                                        f"$delta'<{self.unparser.unparse(funcdef_uast.name)}>({self.unparser.unparse(arg_name)})"
                                    )
                                )
                            ],
                            rhs=[uast.DereferenceExpr(arg_name)],
                        )
                    )
                # case uast.ParamDecl(name=arg_name):
                # pre_process_vertices.append(
                #     uast.Assignment(
                #         "=",
                #         lhs=[
                #             uast.Name(
                #                 f"$delta<{self.unparser.unparse(call_expr.name)}>({self.unparser.unparse(arg_name)})"
                #             )
                #         ],
                #         rhs=[actual_param],
                #     )
                # )
                # pass
                case _:
                    raise NotImplementedError(formal_param)

        print("=======init======")
        for pre_prep in init_vertices:
            print(self.unparser.unparse(pre_prep))
        print("=======final======")
        for pre_prep in final_vertices:
            print(self.unparser.unparse(pre_prep))
        return (
            init_vertices,
            final_vertices,
        )

    def create_output_nodes(self, call_expr: uast.CallExpr):
        pre_process_vertices = []
        post_process_vertices = []

        init_vertices = []
        final_vertices = []

        funcdef_uast = UASTQuery.get_method(
            self.file_uast, self.unparser.unparse(call_expr.name)
        )

        # 解析参数输入节点
        for formal_param, actual_param in zip(
            funcdef_uast.type.pos_args, call_expr.arguments
        ):
            match formal_param:
                case uast.ParamDecl(type=uast.AddrReferenceType(), name=arg_name):
                    # print(formal_param.type, arg_name)
                    pre_process_vertices.append(
                        uast.Assignment(
                            "=",
                            lhs=[
                                uast.DereferenceExpr(
                                    uast.Name(
                                        f"$delta<{self.unparser.unparse(call_expr.name)}>({self.unparser.unparse(arg_name)})"
                                    )
                                )
                            ],
                            rhs=[uast.DereferenceExpr(actual_param)],
                        )
                    )
                    # 添加后处理节点
                    if not isinstance(actual_param, uast.Literal):
                        post_process_vertices.append(
                            uast.Assignment(
                                "=",
                                lhs=[uast.DereferenceExpr(actual_param)],
                                rhs=[
                                    uast.DereferenceExpr(
                                        uast.Name(
                                            f"$delta'<{self.unparser.unparse(call_expr.name)}>({self.unparser.unparse(arg_name)})"
                                        )
                                    )
                                ],
                            )
                        )
                    init_vertices.append(
                        uast.Assignment(
                            "=",
                            lhs=[uast.DereferenceExpr(arg_name)],
                            rhs=[
                                uast.DereferenceExpr(
                                    uast.Name(
                                        f"$delta<{self.unparser.unparse(call_expr.name)}>({self.unparser.unparse(arg_name)})"
                                    )
                                )
                            ],
                        )
                    )
                    final_vertices.append(
                        uast.Assignment(
                            "=",
                            lhs=[
                                uast.DereferenceExpr(
                                    uast.Name(
                                        f"$delta'<{self.unparser.unparse(call_expr.name)}>({self.unparser.unparse(arg_name)})"
                                    )
                                )
                            ],
                            rhs=[uast.DereferenceExpr(arg_name)],
                        )
                    )
                case uast.ParamDecl(name=arg_name):
                    pre_process_vertices.append(
                        uast.Assignment(
                            "=",
                            lhs=[
                                uast.Name(
                                    f"$delta<{self.unparser.unparse(call_expr.name)}>({self.unparser.unparse(arg_name)})"
                                )
                            ],
                            rhs=[actual_param],
                        )
                    )

        print("=======pre proc=======")
        for pre_prep in pre_process_vertices:
            print(self.unparser.unparse(pre_prep))
        print("=======post proc======")

        for pre_prep in post_process_vertices:
            print(self.unparser.unparse(pre_prep))
        # print("=======init======")
        # for pre_prep in init_vertices:
        #     print(self.unparser.unparse(pre_prep))
        # print("=======final======")
        # for pre_prep in final_vertices:
        #     print(self.unparser.unparse(pre_prep))
        return (
            pre_process_vertices,
            post_process_vertices,
            init_vertices,
            final_vertices,
        )

    def analyse_function(self, func_name: str):
        # 通过UASTQuery，获取函数的UAST
        func_uast = UASTQuery.get_method(self.file_uast, func_name)

        # 抽取代码属性图
        cpgs = get_method_cpg(func_uast)

        pdg_main: nx.DiGraph = cpgs.pdg_nx.reverse(False)
        # fm.dot_dump(pdg_main, "p/"cdg_main.dot", export_png=True)
        init_vertices, final_vertices = self.create_input_nodes(func_uast)
        cfg = extract_cfg_from_method(func_uast, remove_empty_nodes=False)

        if func_name != "main":
            # 在CFG的头部和尾部添加初始和结束节点，用于参数传递
            # 除了主函数之外，都需要添加。
            add_nodes_on_cfg_head_or_tail(cfg, init_vertices, "head")
            add_nodes_on_cfg_head_or_tail(cfg, final_vertices, "tail")

        edges_to_ignore_in_ddg = []
        for node_id in pdg_main.nodes:
            stmts = cpgs._cfg.get_block(node_id).statements
            if len(stmts) == 0:
                continue

            if len(exprs := (stmts[0].filter_by(uast.CallExpr).l)) > 0:
                assert len(exprs) == 1
                call_expr = exprs[0]

                # 创建参数输入输出节点
                # preprocess 和 postprocess 要添加到当前的函数上
                preprocess, postprocess, init_vertices, final_vertices = (
                    self.create_output_nodes(call_expr)
                )

                fm.dot_dump(cfg.to_networkx(), f"cfg_{func_name}.dot", export_png=True)

                ctrl_edge = add_controlled_stmt_nodes(
                    cfg, node_id, preprocess + postprocess
                )
                edges_to_ignore_in_ddg.append(ctrl_edge)
                # # CDG的算法只适用于有向无环图？
                # cdg_nx = build_control_dependence_graph(cfg.topology)
                # add_label(cfg.to_networkx(), cdg_nx)
                # fm.dot_dump(
                #     cfg.to_networkx(), f"cfg_{func_name}_processed.dot", export_png=True
                # )
                # # ddg_nx = ddg
                # fm.dot_dump(cdg_nx, f"cdg_{func_name}_processed.dot", export_png=True)

        cfg_nx = cfg.to_networkx()
        cdg_nx = get_cdg_topology(cfg)
        cfg_for_ddg = deepcopy(cfg)
        for ctrl_edge in edges_to_ignore_in_ddg:
            cfg_for_ddg.remove_edge(ctrl_edge[0], ctrl_edge[1])
        new_ddg = get_ddg_topology(
            cfg_for_ddg,
            [],
        )
        set_edge_label(cdg_nx, style="bold", splines="line")
        set_edge_label(new_ddg, splines="curved")
        new_pdg = compose_pdg_topology(cdg_nx, new_ddg, True)
        add_label(cfg_nx, new_pdg)
        add_label(cfg_nx, new_ddg)
        add_label(cfg_nx, cdg_nx)
        # set_edge_label(new_pdg, )

        fm.dot_dump(new_ddg, f"ddg_{func_name}_processed.dot", export_png=True)
        fm.dot_dump(new_pdg, f"pdg_{func_name}_processed.dot", export_png=True)
        fm.dot_dump(cdg_nx, f"cdg_{func_name}_processed.dot", export_png=True)
        return FuncInfo(cfg, new_pdg)


def add_nodes_on_cfg_head_or_tail(
    cfg: uast.CFG,
    statements: list[uast.SourceElement],
    pos: Literal["head", "tail"],
):
    entry_block_id = cfg.entry_block.block_id
    exit_block_id = cfg.exit_block.block_id
    added_blocks: list[uast.BasicBlockData] = []
    for stmt in statements:
        stmt_block = uast.BasicBlockData(cfg.new_node_id(), [stmt], "normal")
        cfg.add_block(stmt_block)
        if len(added_blocks) > 0:
            cfg.add_edge(added_blocks[-1].block_id, stmt_block.block_id)
        added_blocks.append(stmt_block)
    if pos == "head":
        cfg.topology.add_edge(added_blocks[-1].block_id, entry_block_id)
    else:
        cfg.topology.add_edge(exit_block_id, added_blocks[0].block_id)


def add_controlled_stmt_nodes(
    cfg: uast.CFG, control_node_id: int, statements: list[uast.SourceElement]
) -> tuple[int, int]:
    # 控制节点的后继节点
    control_node_successors = list(cfg.topology.successors(control_node_id))
    assert len(control_node_successors) == 1
    original_edge = (control_node_id, control_node_successors[0])
    # 相当于可以增加一个从控制节点，经过statements形成的节点链，指向后继节点的分支，而原连边保持不动。
    # 这样会使得控制流图与代码的实际逻辑不符，但可以增加控制依赖

    added_blocks: list[uast.BasicBlockData] = []
    for stmt in statements:
        stmt_block = uast.BasicBlockData(cfg.new_node_id(), [stmt], "normal")
        cfg.add_block(stmt_block)
        if len(added_blocks) > 0:
            cfg.add_edge(added_blocks[-1].block_id, stmt_block.block_id)
        added_blocks.append(stmt_block)

    cfg.topology.add_edge(control_node_id, added_blocks[0].block_id)
    cfg.topology.add_edge(added_blocks[-1].block_id, control_node_successors[0])
    return original_edge


@dataclass
class FuncInfo:
    cfg: uast.CFG
    pdg: nx.DiGraph
    input_nodes: dict[str, int] = field(default_factory=dict)
    output_nodes: dict[str, int] = field(default_factory=dict)


func_info: dict[str, FuncInfo] = {}

analyzer = Analyzer()
main_info = analyzer.analyse_function("main")
a_info = analyzer.analyse_function("A")
analyzer.analyse_function("Increment")
analyzer.analyse_function("Add")

# cpg_A = get_method_cpg(A_uast)
# pdg_A = cpg_A.pdg_nx.reverse(True)
# fm.dot_dump(pdg_A, "pdg_A.dot", export_png=True)
